# coding: utf-8
import os
import sys
from cv2 import cv2
import numpy as np
import xml.etree.ElementTree as ET
import shutil

# 各种蓝色的HSV值
blues = [
    [120, 255, 77],
    [120, 255, 128],
    [120, 255, 179],
    [120, 255, 230],
    [120, 229, 255]]
path = sys.path[0]  # 获取当前代码文件的路径
print(path)
jpg_list = [x for x in os.listdir(path) if x.endswith(".jpg")]
xml_list = [x.split('.')[0] for x in os.listdir(path) if x.endswith(".xml")]

for num, k in enumerate(blues):
    for idx1, ix in enumerate(jpg_list):  # 所有图片
        print(ix)
        i1 = ix.split('.')
        i2 = i1[0]
        xml_path = path + '\\' + i2 + '.xml'
        tree = ET.parse(xml_path)
        rect = {}
        root = tree.getroot()
        img = cv2.imread(ix)
        for idx, ob in enumerate(root.iter('object')):  # 一张图片中所有主体，一般只有一个
            for bnd in ob.iter('bndbox'):  # 一个主体中所有box
                for x1 in bnd.iter('xmin'):
                    rect['x1_{}'.format(idx)] = int(x1.text)
                for y1 in bnd.iter('ymin'):
                    rect['y1_{}'.format(idx)] = int(y1.text)
                for x2 in bnd.iter('xmax'):
                    rect['x2_{}'.format(idx)] = int(x2.text)
                for y2 in bnd.iter('ymax'):
                    rect['y2_{}'.format(idx)] = int(y2.text)

                cropImg = img[rect['y1_{}'.format(idx)]:rect['y2_{}'.format(idx)],
                          rect['x1_{}'.format(idx)]:rect['x2_{}'.format(idx)]]
                '''try用来判断cropImg 里面是否有值'''
                try:
                    hsv = cv2.cvtColor(cropImg, cv2.COLOR_BGR2HSV)
                except:
                    print('this bndbox is none')
                else:
                    HSV = hsv.copy()
                    for i in range(HSV.shape[0]):
                        for j in range(HSV.shape[1]):
                            '''判断蓝色'''
                            if (HSV[i, j, 0] > 100 and HSV[i, j, 0] < 124 and
                                    HSV[i, j, 1] > 43 and HSV[i, j, 1] < 255 and
                                    HSV[i, j, 2] > 46 and HSV[i, j, 2] < 255):
                                # '''判断红色'''
                                # if ((HSV[i,j,0]>0 and HSV[i,j,0]<10 or HSV[i,j,0]>156 and HSV[i,j,0]<180) and
                                #     HSV[i,j,1]>43 and HSV[i,j,1]<255 and
                                #     HSV[i,j,2]>46 and HSV[i,j,2]<255):
                                print(HSV[i, j, :])
                                HSV[i, j, :] = k
                    img1 = cv2.cvtColor(HSV, cv2.COLOR_HSV2BGR)
                    img[rect['y1_{}'.format(idx)]:rect['y2_{}'.format(idx)],
                    rect['x1_{}'.format(idx)]:rect['x2_{}'.format(idx)]] = img1

        cv2.imwrite(str(i2) + '_' + str(num) + ".jpg", img)
'''生成每张图片对应的XML文件'''
for num, k in enumerate(blues):
    for xml_id, xml_name in enumerate(xml_list):
        xml_file = xml_name + '.xml'
        des_xml = xml_name + '_' + str(num) + '.xml'
        xml_path = os.path.join(path, xml_file)
        despath = os.path.join(path, des_xml)
        shutil.copy(xml_path, despath)
